import numpy as np
import torch
import torch.nn.functional as F

from attack.base_attack import BaseAttack

class Attack(BaseAttack):
    def __init__(self, model, config):
        super().__init__('PGD', model, config, batch_size = 100, targeted = False, llc = False)

    def attack_batch(self, images, labels):
        
        adv_images = np.copy(images)
        adv_images = adv_images + np.random.uniform(-self.config['epsilon'], self.config['epsilon'], images.shape).astype('float32') #随机变化（和BIM的区别）

        for num in range(self.config['num_steps']):
            var_images = torch.from_numpy(adv_images).to(self.model.device)
            var_images.requires_grad = True
            var_labels = torch.LongTensor(labels).to(self.model.device)

            self.model.eval()
            output = self.model(var_images)
            loss = F.cross_entropy(output, var_labels) #逆序排列
            loss.backward()
            grad_sign = var_images.grad.sign().cpu().numpy()

            adv_images = adv_images + self.config['eps_iter'] * grad_sign
            adv_images = np.clip(adv_images, images - self.config['epsilon'], images + self.config['epsilon'])
            adv_images = np.clip(adv_images, 0.0, 1.0)

        return adv_images
